package com.zxs.xss.servlet;

import com.zxs.xss.util.HtmlFilterKit;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import org.springframework.util.StreamUtils;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.Map;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
	private static final String APPLICATION_JSON_VALUE = "application/json";
	private static final String  CHARSET = "UTF-8";

	private final HttpServletRequest orgRequest;
	private final HtmlFilterKit htmlFilter;
	private byte[] cachedBody;

	public XssHttpServletRequestWrapper(HttpServletRequest request, HtmlFilterKit htmlFilter) throws IOException {
		super(request);
		this.htmlFilter = htmlFilter;
		this.orgRequest = request;
		if (APPLICATION_JSON_VALUE.equalsIgnoreCase(request.getContentType())) {
			cachedBody = StreamUtils.copyToByteArray(request.getInputStream());
		}
	}

	@Override
	public ServletInputStream getInputStream() throws IOException {
		if (cachedBody != null) {
			String str = new String(cachedBody,CHARSET);
			String json = xssEncode(str);
			return new CachedServletInputStream(json.getBytes(CHARSET));
		}
		return super.getInputStream();
	}

	@Override
	public String getParameter(String name) {
		String value = super.getParameter(name);
		return StringUtils.isEmpty(value) ? value : xssEncode(value);
	}

	@Override
	public String[] getParameterValues(String name) {
		String[] originalValues = super.getParameterValues(name);

		if (originalValues == null || originalValues.length == 0) {
			return originalValues;
		}

		return Arrays.stream(originalValues)
				.filter(value -> !StringUtils.isEmpty(value))
				.map(this::xssEncode)
				.toArray(String[]::new);
	}

	@Override
	public Map<String, String[]> getParameterMap() {
		Map<String, String[]> originalParameters = super.getParameterMap();

		if (originalParameters == null || originalParameters.isEmpty()) {
			return originalParameters;
		}

		Map<String, String[]> filteredParameters = new LinkedHashMap<>(originalParameters.size());

		for (Map.Entry<String, String[]> entry : originalParameters.entrySet()) {
			String[] originalValues = entry.getValue();
			if (originalValues != null && originalValues.length > 0) {
				// Create a new array to avoid modifying the original values
				String[] encodedValues = Arrays.stream(originalValues)
						.filter(value -> !StringUtils.isEmpty(value))
						.map(this::xssEncode)
						.toArray(String[]::new);
				filteredParameters.put(entry.getKey(), encodedValues);
			}
		}

		return filteredParameters;
	}

	@Override
	public String getHeader(String name) {
		String value = super.getHeader(name);
		return StringUtils.isEmpty(value) ? value : xssEncode(value);
	}

	private String xssEncode(String input) {
		return htmlFilter.filter(input);
	}

	public HttpServletRequest getOrgRequest() {
		return orgRequest;
	}

	public static HttpServletRequest getOrgRequest(HttpServletRequest request) {
		if (request instanceof XssHttpServletRequestWrapper) {
			return ((XssHttpServletRequestWrapper) request).getOrgRequest();
		}
		return request;
	}
}
